# -*- coding: utf-8 -*-
from gevent import monkey
monkey.patch_all(Event=True,subprocess=True)
import zmq.green as zmq
import gevent
from gevent import queue as Queue
from gevent.queue import Empty
from gevent import Greenlet
from random import randint
import time

from alphabet import *

ctx = zmq.Context()


class Publish(object):

    def __init__(self, socket):
        self._sock = socket
        self.task = Queue.Queue()
        self._started = False
        self.run = Greenlet.spawn(self.listen)

    def send(self,command):
        self.task.put(command)

    def listen(self):
        while 1:
            command = self.task.get()
            self._sock.send_multipart(command)

    def start(self):
        if self._started:
            raise RuntimeError("Publish already start.")
        self.run.start()
        self._started = True


class Route(object):

    def __init__(self, socket, handle=None):
        """
        :param socket: ZMQ's socket
        :param handle: ROUTE socket handle
                def handle(message):
                    identify,_,payload = message
                    ...
                    return [identify,'',data]
        """
        if not handle:
            raise RuntimeError("Route need handle's method to dispatch request")
        self._sock = socket
        self.task = Queue.Queue()
        self._started = False
        self.handle = handle
        self.recv = Greenlet.spawn(self.listen)
        self.send = Greenlet.spawn(self.task_handle)

    def before_handle(self, message):
        result = self.handle(message)
        if result:
            self.task.put(result)

    def task_handle(self):
        while 1:
            message = self.task.get()
            self._sock.send_multipart(message)

    def listen(self):
        while 1:
            message = self._sock.recv_multipart()
            Greenlet.spawn(self.before_handle, message).start()

    def start(self):
        if self._started:
            raise RuntimeError("Route already start.")
        self.recv.start()
        self.send.start()
        self._started = True


class Broker(object):

    def __init__(self, check_period=WORKERTIME):
        request_sock = ctx.socket(zmq.ROUTER)
        request_sock.bind("tcp://0.0.0.0:%d"%(REQ_PORT,))
        self.request = Route(request_sock, self.request_handle)

        dispose_sock = ctx.socket(zmq.ROUTER)
        dispose_sock.bind("tcp://0.0.0.0:%d"%(RES_PORT,))
        self.dispose = Route(dispose_sock, self.dispose_handle)

        publish_sock = ctx.socket(zmq.PUB)
        publish_sock.bind("tcp://0.0.0.0:%d"%(CHA_PORT,))
        self.publish = Publish(publish_sock)

        self.workers = {}
        self.check_period = check_period
        self.client_result = {}

        self.request_message = {
            FINDONE:self._request_FINDONE,
            SENDBLOCK:self._request_SENDBLOCK,
            SENDNOBLOCK:self._request_SENDNOBLOCK,
            GETLIST:self._request_GETLIST,
            GETTOTAL:self._request_GETTOTAL,
        }
        self.dispose_message = {
            FINDONE:self._dispose_FINDONE,
            SENDBLOCK:self._dispose_SENDBLOCK,
            SENDNOBLOCK:self._dispose_SENDNOBLOCK,
            HEARTBEAT:self._dispose_HEARTBEAT,
            GETLIST:self._dispose_GETLIST,
            GETTOTAL:self._dispose_GETTOTAL,
        }

    def _request_FINDONE(self, box):
        result = TIMEOUT
        for i in self.workers.keys():
            try:
                result = box.get(timeout=FINDTIMEOUT)
                if result == OK:
                    return result
            except Empty:
                pass
        return result

    def _request_SENDBLOCK(self, box):
        result = TIMEOUT
        for i in self.workers.keys():
            try:
                result = box.get(timeout=DEMANDTIMEOUT)
                # 超时时间与Demand的超时时间相同
                if result == OK:
                    return result
            except Empty:
                pass
        return result

    def _request_SENDNOBLOCK(self, box):
        result = TIMEOUT
        for i in self.workers.keys():
            try:
                result = box.get(timeout=FINDTIMEOUT)
                # 超时时间与查找的超时时间相同
                if result == OK:
                    return result
            except Empty:
                pass
        return result

    def _request_GETLIST(self, box):
        result = []
        for i in self.workers.keys():
            try:
                result.extend(box.get(timeout=FINDTIMEOUT))
                # 超时时间与查找的超时时间相同
            except Empty:
                pass
        return result

    def _request_GETTOTAL(self, box):
        result = 0
        for i in self.workers.keys():
            try:
                result += int(box.get(timeout=FINDTIMEOUT))
                # 超时时间与查找的超时时间相同
            except Empty:
                pass
        return str(result)

    def request_handle(self, message):
        """
        message is [client_id, '', COMMAND, ...(extra data)]
        """
        # print 'request:',message
        client_id = message[0]
        instruct = message[2]
        box = Queue.Queue()
        client_ident = '%s%d'%(client_id, randint(0, 99999999))
        self.client_result[client_ident] = box
        message[0] = client_ident
        self.publish.send(message)
        result = self.request_message[instruct](box)
        del self.client_result[client_ident]
        if isinstance(result,list):
            request = [client_id, '']
            if result:
                request[2:] = result
            else:
                request.append('')
            return request
        else:
            return [client_id, '', result]

    def _dispose_HEARTBEAT(self, worker_id, data):
        self.workers[worker_id] = time.time()

    def _dispose_FINDONE(self, worker_id, data):
        try:
            client_ident = data[0]
            box = self.client_result[client_ident]
            box.put(data[1])
        except KeyError:
            return None

    def _dispose_SENDBLOCK(self, worker_id, data):
        try:
            client_ident = data[0]
            box = self.client_result[client_ident]
            box.put(data[1])
        except KeyError:
            return None

    def _dispose_SENDNOBLOCK(self, worker_id, data):
        try:
            client_ident = data[0]
            box = self.client_result[client_ident]
            box.put(data[1])
        except KeyError:
            return None

    def _dispose_GETLIST(self, worker_id, data):
        try:
            client_ident = data[0]
            box = self.client_result[client_ident]
            box.put(data[1:])
        except KeyError:
            return None

    def _dispose_GETTOTAL(self, worker_id, data):
        try:
            client_ident = data[0]
            box = self.client_result[client_ident]
            box.put(data[1])
        except KeyError:
            return None

    def dispose_handle(self, message):
        """
        message is [worker_id, '', COMMAND, ...(extra data)]
        """
        # print 'dispose:',message
        worker_id = message[0]
        instruct = message[2]
        data = message[3:]
        self.dispose_message[instruct](worker_id, data)
        return [worker_id, '', 'OK']

    def start(self):
        self.request.start()
        self.dispose.start()

    def run(self):
        self.start()
        while 1:
            gevent.sleep(self.check_period)
            now = time.time()
            for w in self.workers.keys():
                if now - self.workers[w] >= self.check_period:
                    del self.workers[w]



if __name__ == "__main__":
    broker = Broker()
    broker.run()



